package com.aaron.loadbalance.loadbalancer;

import com.aaron.loadbalance.AbstractLoadBalance;
import com.aaron.remoting.dto.RpcRequest;
import lombok.extern.slf4j.Slf4j;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 基于 hash 一致性的 ConsistentHashLoadBalance。
 * 参考http://tianxiaobo.com/2018/11/29/Dubbo-%E6%BA%90%E7%A0%81%E5%88%86%E6%9E%90-%E9%9B%86%E7%BE%A4%E5%AE%B9%E9%94%99%E4%B9%8B-LoadBalance/
 */
@Slf4j
public class ConsistentHashLoadBalance extends AbstractLoadBalance {
    // key:String类型的rpcServiceName，value:内部类ConsistentHashSelector类型的
    private final ConcurrentHashMap<String, ConsistentHashSelector> selectors = new ConcurrentHashMap<>();

    // 从保存远程服务地址的list集合中调用一个serviceAddresses
    @Override
    protected String doSelect(List<String> serviceAddresses, RpcRequest rpcRequest) {
        // 获取serviceAddresses原始的hashCode
        int identityHashCode = System.identityHashCode(serviceAddresses);
        // ConsistentHashSelector是一个内部类，相当于缓存的效果
        String rpcServiceName = rpcRequest.getRpcServiceName();
        ConsistentHashSelector selector = selectors.get(rpcServiceName);

        /**
         * 如果 serviceAddresses 是一个新的 List 对象，意味着服务提供者数量发生了变化，可能新增也可能减少了。
         * 此时 selector.identityHashCode != identityHashCode 条件成立
         * identityHashCode是 旧List<String> serviceAddresses对象的identityHashCode
         * selector.identityHashCode是 ConsistentHashSelector类的identityHashCode
         */
        if (selector == null || selector.identityHashCode != identityHashCode) {
            // 创建新的ConsistentHashSelector内部类缓存
            selectors.put(rpcServiceName, new ConsistentHashSelector(serviceAddresses, 160, identityHashCode));
            selector = selectors.get(rpcServiceName);
        }

        // 调用ConsistentHashSelector 的 select 方法选择serviceAddresses
        return selector.select(rpcServiceName + Arrays.stream(rpcRequest.getParameters()));
    }


    // 内部类，保存虚拟节点的集合，以及该集合的identityHashCode
    static class ConsistentHashSelector {
        // 使用 TreeMap 存储 serviceAddresses 虚拟节点，key就是hash值，value是地址
        private final TreeMap<Long, String> virtualInvokers;

        // 就是List<String> serviceAddresses对象的identityHashCode
        private final int identityHashCode;

        // 内部类的构造方法
        ConsistentHashSelector(List<String> invokers, int replicaNumber, int identityHashCode) {
            this.virtualInvokers = new TreeMap<>();
            this.identityHashCode = identityHashCode;
            // invokers是String类型的serviceAddresses，一个serviceAddresses在virtualInvokers这个map中就有160个节点
            for (String invoker : invokers) {
                // 按照前面的传参replicaNumber=160,那么就是进行了40次运算
                for (int i = 0; i < replicaNumber / 4; i++) {
                    // 对invoker+i进行md5运算，得到一个长度为16个字节数组
                    byte[] digest = md5(invoker + i);
                    // 对digest部分字节进行4次hash运算，得到四个不同的long型正整数
                    for (int h = 0; h < 4; h++) {
                        // h = 0 时，取 digest 中下标为 0 ~ 3 的4个字节进行位运算
                        // h = 1 时，取 digest 中下标为 4 ~ 7 的4个字节进行位运算
                        // h = 2 时，取 digest 中下标为 8 ~ 11 的4个字节进行位运算
                        // h = 3 时，取 digest 中下标为 12 ~ 15 的4个字节进行位运算
                        long m = hash(digest, h);
                        // 将 hash 到 invoker 的映射关系存储到 virtualInvokers 中，
                        // virtualInvokers 中的元素要有序，因此选用 TreeMap 作为存储结构
                        // 得到了160个虚拟节点，存储到了virtualInvokers这个TreeMap中
                        virtualInvokers.put(m, invoker); // invoker就是一个serviceAddresses
                        // virtualInvokers这个TreeMap中  应该有serviceAddresses.size * 160个节点
                    }
                }
            }
        }

        // rpcServiceKey就是 rpcServiceName + Arrays.stream(rpcRequest.getParameters())
        public String select(String rpcServiceKey) {
            // 对参数 key 进行 md5 运算
            byte[] digest = md5(rpcServiceKey);
            // 取 digest 数组的前四个字节进行 hash 运算，为什么只是前4个字节？之前还虚拟出了其他字节的节点？
            // 再将hash值传给selectForKey方法，得到服务地址
            long hashCode1 = hash(digest, 0);
            return selectForKey(hashCode1);
        }

        public String selectForKey(long hashCode) {
            /**
             * tailMap()方法与firstEntry()方法表示：
             *      到virtualInvokers这个TreeMap中查找第一个节点值大于或等于当前hashCode的节点entry
             *      virtualInvokers中第一个大于或等于hashCode的节点entry，对应的数据结构是<Long, String>
             */
            Map.Entry<Long, String> entry = virtualInvokers.tailMap(hashCode, true).firstEntry();
            // 如果hashCode 大于 圆环上最大的位置，此时节点entry = null，需要将 TreeMap 的头结点赋值给 entry
            if (entry == null) {
                entry = virtualInvokers.firstEntry();
            }
            // 返回的是String类型的serviceAddresses。可以看出地址是保存在虚拟节点里面的
            return entry.getValue();
        }

        // 只要知道它是在做运算，返回的是字节数组。模仿dubbo里面的
        static byte[] md5(String key) {
            MessageDigest md;
            try {
                md = MessageDigest.getInstance("MD5");
                byte[] bytes = key.getBytes(StandardCharsets.UTF_8);
                md.update(bytes);
            } catch (NoSuchAlgorithmException e) {
                throw new IllegalStateException(e.getMessage(), e);
            }
            return md.digest();
        }

        // 只要知道它是在做运算。dubbo里面的
        static long hash(byte[] digest, int idx) {
            return ((long) (digest[3 + idx * 4] & 255) << 24 | (long) (digest[2 + idx * 4] & 255) << 16 | (long) (digest[1 + idx * 4] & 255) << 8 | (long) (digest[idx * 4] & 255)) & 4294967295L;
        }
    }
}
